import tensorflow as tf
import numpy

def block2(input_tensor,out_dim):
    conv = tf.keras.layers.Conv2D(out_dim//4,kernel_size=1,padding="SAME",activation=tf.nn.relu)(input_tensor)
    conv = tf.keras.layers.BatchNormalization()(conv)
    conv = tf.keras.layers.Conv2D(out_dim//4,kernel_size=3,padding="SAME",activation=tf.nn.relu)(conv)
    conv = tf.keras.layers.BatchNormalization()(conv)
    conv = tf.keras.layers.Conv2D(out_dim,kernel_size=1,padding="SAME",activation=tf.nn.relu)(conv)
    out = tf.keras.layers.Add()([input_tensor,conv])
    out = tf.keras.layers.ReLU()(out)
    return out


def Model(n_dim = 5):
    input_xs = tf.keras.Input([13,14,100])

    out_dim = 32
    conv0 = tf.keras.layers.Conv2D(out_dim,kernel_size=3,padding="SAME",activation=tf.nn.relu)(input_xs)
    conv0 = tf.keras.layers.BatchNormalization()(conv0)
    conv0 = tf.keras.layers.AveragePooling2D(pool_size=[1,1],padding="SAME")(conv0)


    conv1 = tf.keras.layers.Conv2D(out_dim//4,kernel_size=3,padding="SAME",activation=tf.nn.relu)(conv0)
    conv1 = tf.keras.layers.BatchNormalization()(conv1)
    conv1 = tf.keras.layers.Conv2D(out_dim,kernel_size=1,padding="SAME",activation=tf.nn.relu)(conv1)

    conv0 = tf.keras.layers.Conv2D(out_dim,kernel_size=1,padding="SAME",activation=tf.nn.relu)(input_xs)
    conv0 = tf.keras.layers.BatchNormalization()(conv0)

    out = tf.keras.layers.Add()([conv0,conv1])
    out = tf.keras.layers.ReLU()(out)

    out = block2(out,out_dim)

    out = block2(out,out_dim)

    out = tf.keras.layers.AveragePooling2D(pool_size=[2,2],padding="SAME")(out)
    out = tf.keras.layers.Dense(128,activation=tf.nn.relu)(out)
    out = tf.keras.layers.BatchNormalization()(out)
    logits = tf.keras.layers.Dense(n_dim,activation=tf.nn.softmax)(out)


    print(logits)
    model = tf.keras.Model(inputs=input_xs,outputs=logits)


    return model

if __name__ == "__main__":
    Model()
    model = Model()
    print(model.summary())
